package com.feidee.fd.sml.algorithm.ml

import com.feidee.fd.sml.algorithm.component.ml.classification.{NaiveBayesComponent, NaiveBayesParam}
import com.feidee.fd.sml.algorithm.forecast.StageFinder
import com.feidee.fd.sml.algorithm.util.{TestingDataGenerator, ToolClass}
import org.apache.spark.ml.classification.NaiveBayesModel
import org.scalatest.FunSuite

/**
  * @Author songhaicheng
  * @Date 2018/08/28
  * @Email: haicheng_song@sui.com
  */
class NaiveBayesComponentSuite extends FunSuite {

  val nb = new NaiveBayesComponent
  val paramStr: String =
    """
      |{
      | 'rawPredictionCol': 'raw',
      |	'modelPath': '',
      |	'metrics': ['auc', 'pr', 'nothing', 'accuracy'],
      | 'modelType': 'multinomial'
      |}
    """.stripMargin
  val param: NaiveBayesParam = nb.parseParam(new ToolClass().encrypt(paramStr))

  test("NaiveBayes parameter construction test") {
    assert(param.input_pt == null & param.output_pt == null & "features".equals(param.featuresCol) & "label".equals(param.labelCol) &
      "prediction".equals(param.predictionCol) & "".equals(param.modelPath) & param.metrics.sameElements(Array("auc", "pr", "nothing", "accuracy")) &
      "raw".equals(param.rawPredictionCol) & "probability".equals(param.probabilityCol) & param.weightCol == null &
      param.smoothing == 1.0 & "multinomial".equals(param.modelType) & param.thresholds.length == 0
    )
  }

  test("NaiveBayes training function test") {
    val testingData = TestingDataGenerator.makeOrderedTrainingData(5, 20)

    val model = nb.train(param, testingData)
    assert(new StageFinder[NaiveBayesModel](model).findWithOrder().get.numClasses == 2)
    assert(model.transform(testingData).count() == 20)

  }

}
